package com.cryptoneedle.center.source.util;


import com.cryptoneedle.center.source.constant.Constant;
import com.cryptoneedle.center.source.entity.metadata.SourceServer;
import com.cryptoneedle.center.source.properties.SshProperties;
import com.cryptoneedle.center.source.service.metadata.SourceServerService;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import jakarta.annotation.Resource;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * <p>description: SSH隧道服务 </p>
 *
 * @author CryptoNeedle
 * @date 2025-06-23
 */
@Slf4j
@Service
public class SshManager implements InitializingBean {
    
    @Resource
    private SshProperties sshProperties;
    
    @Resource
    private SourceServerService sourceServerService;
    
    private AtomicInteger openPortIndex;
    
    private Session session;
    
    // eg. key:host:port  value:openPortIndex
    @Getter
    private static final Map<String, Integer> successForwardMap = new ConcurrentHashMap<>();
    
    // eg. key:host  value:port
    @Getter
    private static final Map<String, Integer> failForwardMap = new ConcurrentHashMap<>();
    
    /**
     * 校验Session连接
     */
    @Override
    public void afterPropertiesSet() throws Exception {
        if (!sshProperties.getEnabled()) {
            return;
        }
        
        // 代理
        session = proxy();
        String message;
        if (session == null) {
            message = String.format("[SSH] %s:%s => 会话为空", sshProperties.getHost(), sshProperties.getPort());
            log.error(message);
            throw new RuntimeException(message);
        }
        if (!session.isConnected()) {
            message = String.format("[SSH] %s:%s => 连接失败", sshProperties.getHost(), sshProperties.getPort());
            log.error(message);
            throw new RuntimeException(message);
        }
        
        openPortIndex = new AtomicInteger(sshProperties.getBeginOpenPort());
        
        // 获取需要SSH转发的主机和端口
        List<SourceServer> sourceServers = sourceServerService.findAll();
        // eg. key: 127.0.0.1  value: {3306,5432}
        Map<String, List<Integer>> serverHostPortMap = sourceServers.stream()
                                                                    .collect(Collectors.groupingBy(SourceServer::getHost, Collectors.mapping(SourceServer::getPort, Collectors.toList())));
        
        for (Map.Entry<String, List<Integer>> serverHostPortEntry : serverHostPortMap.entrySet()) {
            String targetHost = serverHostPortEntry.getKey();
            List<Integer> targetPorts = serverHostPortEntry.getValue();
            for (Integer targetPort : targetPorts) {
                forward(targetHost, targetPort);
            }
        }
        
        if (!successForwardMap.isEmpty()) {
            successForwardMap.forEach((ip, targetPort) -> {
                log.info("[SSH] {}:{} -> 转发 => {} -> {}", sshProperties.getHost(), sshProperties.getPort(), ip, targetPort);
            });
        }
        if (!failForwardMap.isEmpty()) {
            failForwardMap.forEach((host, port) -> {
                log.info("[SSH] {}:{} -> 转发失败 => {} -> {}", sshProperties.getHost(), sshProperties.getPort(), host, port);
            });
        }
    }
    
    /**
     * SSH远程连接
     */
    private Session proxy() {
        Session session = null;
        try {
            JSch jSch = new JSch();
            
            // 创建会话
            session = jSch.getSession(sshProperties.getUser(), sshProperties.getHost(), sshProperties.getPort());
            session.setPassword(sshProperties.getPassword());
            session.setConfig("StrictHostKeyChecking", "no");
            session.setTimeout(sshProperties.getTimeout());
            session.connect();
            
            log.info("[SSH] {}:{} -> 连接", sshProperties.getHost(), sshProperties.getPort());
        } catch (JSchException e) {
            log.error("[SSH] {}:{} -> 连接 -> 异常 => {}", sshProperties.getHost(), sshProperties.getPort(), e.getMessage());
        }
        return session;
    }
    
    /**
     * SSH隧道代理
     */
    
    public synchronized boolean forward(String targetHost, Integer targetPort) {
        boolean successFlag = forward(targetHost, targetPort, 0);
        if (successFlag) {
            successForwardMap.put(targetHost + ":" + targetPort, openPortIndex.get() - 1);
        } else {
            failForwardMap.put(targetHost, targetPort);
        }
        return successFlag;
    }
    
    private synchronized boolean forward(String targetHost, Integer targetPort, int errorNum) {
        try {
            session.setPortForwardingL(Constant.LOCAL_HOST, openPortIndexIncrement(), targetHost, targetPort);
            return true;
        } catch (JSchException e) {
            
            log.debug("[SSH] {}:{} -> 转发 -> 异常 => {}:{} -> {}:{}, 第{}次尝试 => {}", sshProperties.getHost(), sshProperties.getPort(), targetHost, targetPort, Constant.LOCAL_HOST, openPortIndex.get() - 1, errorNum + 1, e.getMessage());
            
            // 允许尝试十次
            if (errorNum <= 10) {
                return forward(targetHost, targetPort, errorNum + 1);
            }
            return false;
        }
    }
    
    private synchronized Integer openPortIndexIncrement() {
        if (openPortIndex.get() + 1 >= 65535) {
            log.error("[SSH] {}:{} -> 可用端口耗尽，将进行重置", sshProperties.getHost(), sshProperties.getPort());
            openPortIndex.set(sshProperties.getBeginOpenPort());
        }
        Integer old = openPortIndex.get();
        openPortIndex.getAndIncrement();
        return old;
    }
    
    /**
     * 工具
     */
    
    public Integer getProxyPort(String host, Integer port) {
        String address = host + ":" + port;
        Integer proxyPort = successForwardMap.get(address);
        if (proxyPort == null) {
            forward(host, port);
            proxyPort = successForwardMap.get(address);
            if (proxyPort == null) {
                throw new RuntimeException("[SSH] " + address + " 代理端口失败");
            }
        }
        return proxyPort;
    }
}